import base64
import hashlib

from flask import Flask, request, jsonify
import requests
import uuid
import datetime
import os
from collections import Counter
from multiprocessing import Pool, Manager
from datetime import datetime
import warnings
from multiprocessing import Process

# 忽略所有警告
warnings.filterwarnings("ignore")
app = Flask(__name__)

# 使用 Manager 创建共享的 tasks 字典
manager = Manager()
tasks = manager.dict()  # 替换原有的普通字典

# 评率阈值（最低评率不能超过 18）
FREQUENCY_THRESHOLD = 18

# 最大并发数（你可以调整）
MAX_PROCESSES = 150

def load_dict_from_file(filename):
    """ 从文件读取用户名或密码字典，按行解析 """
    if not os.path.exists(filename):
        return []

    with open(filename, "r", encoding="utf-8") as file:
        return [line.strip() for line in file if line.strip()]

USERNAME_DICT = load_dict_from_file("mynamedic.txt")
# PASSWORD_DICT = load_dict_from_file("mypassword.txt")
'''下面的测试专用'''
# USERNAME_DICT = load_dict_from_file("testname.txt")
# PASSWORD_DICT = load_dict_from_file("testpassword.txt")

def generate_payload(base_body, username, password , password_type):
    """ 按模式替换用户名和密码，不改变 body 格式 """
    if password_type == "md5":
        return base_body.replace("aaaaaa", username).replace("875f26fdb1cecf20ceb4ca028263dec6", password)
    elif password_type == "base64":
        return base_body.replace("aaaaaa", username).replace("YmJiYmJi", password)
    elif password_type == "plain":
        return base_body.replace("aaaaaa", username).replace("bbbbbb", password)

def generate_get_url(url, username, password, password_type):
    """ 根据用户名和密码修改 GET 请求的 URL 查询参数 """
        # 替换 URL 查询参数中的用户名和密码
    if password_type == "md5":
        return url.replace("aaaaaa", username).replace("875f26fdb1cecf20ceb4ca028263dec6", password)
    elif password_type == "base64":
        return url.replace("aaaaaa", username).replace("YmJiYmJi", password)
    elif password_type == "plain":
        return url.replace("aaaaaa", username).replace("bbbbbb", password)


def send_request(url, method, headers, body):
    """ 发送 HTTP 请求 """
    try:
        if method == "GET":
            r = requests.get(url, headers=headers, timeout=2, verify=False)
        else:
            r = requests.post(url, headers=headers, data=body.encode('utf-8'), timeout=2, verify=False)

        response_length = len(r.text)
        return response_length
    except requests.exceptions.RequestException:
        return None


def analyze_responses(responses):
    """
    分析所有响应的长度分布，找到评率最低的那个，
    若最低评率 ≤ 20，则认为该响应对应的 payload 为爆破成功
    """
    # lengths = [len(r.text) for r in responses if r]  # 获取所有有效响应的长度
    lengths = responses
    if not lengths:
        return None
    counter = Counter(lengths)  # 计算长度频率
    min_freq = min(counter.values())  # 获取最小评率

    # 如果最低评率小于等于 19，找出对应的 payload
    if min_freq <= FREQUENCY_THRESHOLD:
        success_length = min(counter, key=counter.get)  # 获取评率最低的长度
        return success_length
    return None

def execute_task(task_id):
    """ 执行爆破任务（并发版本） """
    task = tasks.get(task_id)
    if not task:
        return

    url = task["url"]
    method = task["method"]
    headers = task["headers"]
    base_body = task["body"]
    password_type = task["password_type"]


    # 根据 password_type 选择密码字典
    if password_type == "md5":
        password_dict = load_dict_from_file("md5_mypassword.txt")
    elif password_type == "base64":
        password_dict = load_dict_from_file("base64_mypassword.txt")
    elif password_type == "plain":
        password_dict = load_dict_from_file("mypassword.txt")  # 自定义字典
    else:
        password_dict = None

    # 如果没有密码可用，则提前终止
    if not password_dict:
        task["status"] = "failed"
        return
    response_map = {}  # 存储响应长度和对应的用户名密码
    # all_responses = []  # 存储所有请求的响应
    # tested_payloads = []  # 记录已测试的 payload 和响应信息

    # 使用共享列表
    tested_payloads = manager.list()
    all_responses = manager.list()
    task["tested_payloads"] = tested_payloads  # 存储 ListProxy
    task["total_payloads"] = len(all_responses)

    with Pool(processes=MAX_PROCESSES) as pool:
        results = []
        for username in USERNAME_DICT:
            for password in password_dict:
                if method == "POST":
                    # 对 POST 请求，生成 body
                    body = generate_payload(base_body, username, password, password_type)
                    response = pool.apply_async(send_request, args=(url, method, headers, body))
                elif method == "GET":
                    # 对 GET 请求，生成带查询参数的 URL
                    modified_url = generate_get_url(url, username, password, password_type)
                    response = pool.apply_async(send_request, args=(modified_url, method, headers, None))
                results.append((response, username, password))

        # 等待所有请求完成并处理响应
        for response, username, password in results:
            response_length = response.get()
            if response_length:
                # response_length = response.text
                if response_length not in response_map:
                    response_map[response_length] = []
                if len(response_map[response_length]) < 20:
                    response_map[response_length].append((username, password))
                    tested_payloads.append({
                        "username": username,
                        "password": password,
                        "response_length": response_length
                    })
                # response_map[response_length] = (username, password)
                all_responses.append(response_length)
                task["total_payloads"] = len(all_responses)
                # 将结果添加到共享列表


    # 分析响应，找到评率最低的爆破成功 Payload
    success_length = analyze_responses(all_responses)

    if success_length:
        success_credentials = response_map[success_length]

        # 检查并收集所有不重复的凭证
        valid_credentials = []
        seen_usernames = set()

        for username, password in success_credentials:
            if username not in seen_usernames:
                valid_credentials.append({
                    "username": username,
                    "password": password
                })
                seen_usernames.add(username)
            else:
                task["status"] = "failed"
                valid_credentials=[]
                break

        if valid_credentials:
            task["status"] = "success"
            task["credentials"] = valid_credentials  # 直接存储凭证列表
        else:
            task["status"] = "failed"
    else:
        task["status"] = "failed"

@app.route('/create_task', methods=['POST'])
def create_task():
    """ 创建任务并自动执行 """
    data = request.json

    # 使用 Manager 创建共享的任务字典
    task_id = str(uuid.uuid4())
    timestamp = datetime.now().isoformat()

    # 创建共享任务字典（包含 start_time 和 end_time）
    shared_task = manager.dict({
        "task_id": task_id,
        "url": data.get("url"),
        "method": data.get("method"),
        "headers": data.get("headers", {}),
        "body": data.get("body", ""),
        "password_type" : data.get("password_type"),
        "timestamp": timestamp,
        "start_time": timestamp,
        "end_time": None,
        "status": "pending",
        "tested_payloads": manager.list(),
        "total_payloads": None
    })
    tasks[task_id] = shared_task

    # 启动任务进程

    process = Process(target=execute_task, args=(task_id,))
    process.start()

    return jsonify({"message": "Task created and started", "task_id": task_id}), 201

@app.route('/tasks', methods=['GET'])
def get_all_tasks():
    task_list = []
    for task_id in tasks:
        task = tasks[task_id]
        task_info = {
            "task_id": task["task_id"],
            "status": task["status"],
            "url": task["url"],
            "method": task["method"],
            "timestamp": task["timestamp"],
            "rps": None,  # 默认值
            "success_credentials": {
                "username_password": task.get("credentials")
            } if task["status"] == "success" else None
        }

        # 计算 RPS
        if task["status"] in ["success", "failed"]:
            if task["start_time"] and task["end_time"]:
                start = datetime.fromisoformat(task["start_time"])
                end = datetime.fromisoformat(task["end_time"])
                duration = (end - start).total_seconds()
                total = len(task["tested_payloads"])
                if duration > 0:
                    task_info["rps"] = round(total / duration, 2)

        # 实时 RPS（任务执行中）
        elif task["status"] == "pending" and task["start_time"]:
            start = datetime.fromisoformat(task["start_time"])
            now = datetime.now()
            duration = (now - start).total_seconds()
            total = int(task["total_payloads"])
            if duration > 0:
                task_info["rps"] = round(total / duration, 2)

        task_list.append(task_info)

    return jsonify(task_list)


@app.route('/task_status/<task_id>', methods=['GET'])
def get_task_status(task_id):
    """ 查询单个任务的爆破详情 """
    task = tasks.get(task_id)
    if not task:
        return jsonify({"error": "Task not found"}), 404

    # 将 ListProxy 转换为普通列表
    tested_payloads = list(task["tested_payloads"])  # 关键转换

    return jsonify({
        "task_id": task["task_id"],
        "status": task["status"],
        "url": task["url"],
        "method": task["method"],
        "timestamp": task["timestamp"],
        "tested_payloads": tested_payloads,  # 使用转换后的普通列表
        "success_credentials": {
            "username_password": task.get("credentials")
        } if task["status"] == "success" else None
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=4646)
